package com.zhanghe.study.springboot.configuration;

import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
import java.io.UnsupportedEncodingException;

/**
 * @author zh
 * @date 2023/1/10 15:05
 */
public class RestTemplateInterceptor implements ClientHttpRequestInterceptor {


    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {

        this.traceRequest((ClientHttpRequest) request, body);
        ClientHttpResponse response = execution.execute(request, body);
        return response;
    }

    private void traceRequest(ClientHttpRequest request, byte[] body) throws UnsupportedEncodingException {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        if (requestAttributes == null) {
            return;
        }
        HttpServletRequest req = ((ServletRequestAttributes) requestAttributes).getRequest();
        HttpHeaders headers = new HttpHeaders();

        String globalTraceId = req.getHeader("global-trace-id");
        if (globalTraceId == null) {
//            headers.add("global-trace-id", TraceThreadLocal.getGlobalTraceId());
        } else {
            headers.add("global-trace-id", globalTraceId);
        }
        String parentTraceId = req.getHeader("local-trace-id");
        if (parentTraceId == null) {
//            headers.add("parent-trace-id", TraceThreadLocal.getLocalTraceId());
        } else {
            headers.add("parent-trace-id", parentTraceId);
        }
        request.getHeaders().putAll(headers);
    }


}
